"""
子通道分析主模块

包含子通道的几何定义和主要的流动计算
"""

import numpy as np
from ..models.fluid_properties import FluidProperties
from ..models.heat_transfer import HeatTransfer
from .. import config as cfg

class SubChannel:
    """子通道类，表示单个子通道"""
    
    def __init__(self, channel_id, neighbors):
        """
        初始化子通道
        
        参数：
        channel_id: int, 子通道编号
        neighbors: list, 相邻子通道编号列表
        """
        self.id = channel_id
        self.neighbors = neighbors
        self.diameter = cfg.CHANNEL_DIAMETER
        self.area = np.pi * (self.diameter/2)**2
        
        # 初始化场变量
        self.alpha_l = np.zeros(cfg.AXIAL_NODES)  # 液相分数
        self.alpha_v = np.zeros(cfg.AXIAL_NODES)  # 气相分数
        self.u_l = np.zeros(cfg.AXIAL_NODES)      # 液相轴向速度
        self.u_v = np.zeros(cfg.AXIAL_NODES)      # 气相轴向速度
        self.v_l = np.zeros(cfg.AXIAL_NODES)      # 液相横向速度
        self.v_v = np.zeros(cfg.AXIAL_NODES)      # 气相横向速度
        self.pressure = np.zeros(cfg.AXIAL_NODES)  # 压力
        self.T_l = np.zeros(cfg.AXIAL_NODES)      # 液相温度
        self.T_v = np.zeros(cfg.AXIAL_NODES)      # 气相温度
        self.T_w = np.zeros(cfg.AXIAL_NODES)      # 壁面温度
        
        # 初始化物性和传热计算器
        self.fluid_props = FluidProperties()
        self.heat_transfer = HeatTransfer(self.diameter)
        
    def initialize_fields(self, inlet_pressure, inlet_temp, mass_flux):
        """
        初始化场变量
        
        参数：
        inlet_pressure: float, 入口压力 (Pa)
        inlet_temp: float, 入口温度 (K)
        mass_flux: float, 质量流量 (kg/m²s)
        """
        # 获取入口处液相物性
        props = self.fluid_props.get_properties(inlet_pressure, inlet_temp, 'liquid')
        if props is None:
            return
            
        # 设置入口条件
        self.alpha_l[0] = 1.0  # 入口处全为液相
        self.alpha_v[0] = 0.0
        self.u_l[0] = mass_flux / props['density']  # 入口速度
        self.u_v[0] = 0.0
        self.pressure[0] = inlet_pressure
        self.T_l[0] = inlet_temp
        self.T_v[0] = inlet_temp
        self.T_w[0] = inlet_temp
        
        # 初始化其他节点
        for i in range(1, cfg.AXIAL_NODES):
            self.alpha_l[i] = self.alpha_l[0]
            self.alpha_v[i] = self.alpha_v[0]
            self.u_l[i] = self.u_l[0]
            self.u_v[i] = self.u_v[0]
            self.pressure[i] = self.pressure[0]
            self.T_l[i] = self.T_l[0]
            self.T_v[i] = self.T_v[0]
            self.T_w[i] = self.T_w[0]
            
    def calculate_cross_flow(self, i, neighbor_channel):
        """
        计算与相邻子通道的横向流动
        
        参数：
        i: int, 轴向位置索引
        neighbor_channel: SubChannel, 相邻子通道对象
        
        返回：
        float, 横向质量流量 (kg/s)
        """
        # 获取连接系数
        connection_key = (min(self.id, neighbor_channel.id), 
                        max(self.id, neighbor_channel.id))
        if connection_key not in cfg.CONNECTION_COEFFICIENTS:
            return 0.0
            
        s_jk = cfg.CONNECTION_COEFFICIENTS[connection_key]
        
        # 计算压力差驱动的横向流动
        dp = self.pressure[i] - neighbor_channel.pressure[i]
        cross_flow = s_jk * dp * 1e-3  # 横向流与压力差成正比
        
        return cross_flow
        
    def calculate_wall_friction(self, i):
        """
        计算壁面摩擦力
        
        参数：
        i: int, 轴向位置索引
        
        返回：
        tuple, (液相摩擦力, 气相摩擦力) (N/m³)
        """
        # 液相摩擦力
        props_l = self.fluid_props.get_properties(self.pressure[i], self.T_l[i], 'liquid')
        if props_l is None:
            return 0.0, 0.0
            
        Re_l = props_l['density'] * abs(self.u_l[i]) * self.diameter / props_l['viscosity']
        if Re_l < 2300:  # 层流
            f_l = 64 / max(Re_l, 1e-6)
        else:  # 湍流
            f_l = 0.316 * max(Re_l, 1e-6)**(-0.25)
            
        # 限制速度以避免数值溢出
        u_limited = np.clip(self.u_l[i], -100, 100)  # 限制速度在合理范围内
        tau_wl = 0.5 * f_l * props_l['density'] * u_limited * abs(u_limited)
        
        # 添加数值稳定性处理
        if not np.isfinite(tau_wl):
            tau_wl = 0.0
        
        F_wl = tau_wl * np.pi * self.diameter / self.area
        
        # 确保壁面摩擦力有界
        F_wl = np.clip(F_wl, -1e6, 1e6)
        
        # 气相摩擦力
        if self.alpha_v[i] > 0:
            props_v = self.fluid_props.get_properties(self.pressure[i], self.T_v[i], 'vapor')
            if props_v is None:
                return F_wl, 0.0
                
            Re_v = props_v['density'] * abs(self.u_v[i]) * self.diameter / props_v['viscosity']
            if Re_v < 2300:
                f_v = 64 / Re_v
            else:
                f_v = 0.316 * Re_v**(-0.25)
                
            tau_wv = 0.5 * f_v * props_v['density'] * self.u_v[i] * abs(self.u_v[i])
            F_wv = tau_wv * np.pi * self.diameter / self.area
        else:
            F_wv = 0.0
            
        return F_wl, F_wv
        
    def calculate_interfacial_force(self, i):
        """
        计算界面剪切力
        
        参数：
        i: int, 轴向位置索引
        
        返回：
        float, 界面剪切力 (N/m³)
        """
        if self.alpha_v[i] <= 0 or self.alpha_v[i] >= 1:
            return 0.0
            
        props_l = self.fluid_props.get_properties(self.pressure[i], self.T_l[i], 'liquid')
        if props_l is None:
            return 0.0
            
        # 计算相对速度
        v_rel = self.u_v[i] - self.u_l[i]
        
        # 计算界面面积密度（假设球形气泡）
        a_i = 6 * self.alpha_v[i] / self.diameter
        
        # 计算界面阻力系数
        C_i = cfg.BUBBLE_DRAG_COEFF
        
        # 计算界面剪切力
        F_i = 0.5 * C_i * props_l['density'] * a_i * v_rel * abs(v_rel)
        
        return F_i
        
    def calculate_heat_transfer(self, i, q_wall):
        """
        计算传热
        
        参数：
        i: int, 轴向位置索引
        q_wall: float, 壁面热流密度 (W/m²)
        
        返回：
        tuple, (液相温度变化率, 气相温度变化率) (K/s)
        """
        # 获取饱和温度
        T_sat = self.fluid_props.get_saturation_temperature(self.pressure[i])
        if T_sat is None:
            return 0.0, 0.0
            
        # 计算液相传热
        if self.alpha_l[i] > 0:
            if self.T_l[i] < T_sat and self.T_w[i] > T_sat:
                # 过冷核态沸腾
                h_l = self.heat_transfer.calculate_subcooled_boiling_htc(
                    self.T_w[i], T_sat, self.T_l[i], self.pressure[i])
            else:
                # 单相对流
                h_l = self.heat_transfer.calculate_single_phase_htc(
                    self.pressure[i], self.T_l[i], self.u_l[i], 'liquid')
                    
            q_l = h_l * (self.T_w[i] - self.T_l[i])
            props_l = self.fluid_props.get_properties(self.pressure[i], self.T_l[i], 'liquid')
            if props_l is not None:
                dT_l_dt = q_l / (props_l['density'] * props_l['cp'])
            else:
                dT_l_dt = 0.0
        else:
            dT_l_dt = 0.0
            
        # 计算气相传热
        if self.alpha_v[i] > 0:
            h_v = self.heat_transfer.calculate_single_phase_htc(
                self.pressure[i], self.T_v[i], self.u_v[i], 'vapor')
            q_v = h_v * (self.T_w[i] - self.T_v[i])
            props_v = self.fluid_props.get_properties(self.pressure[i], self.T_v[i], 'vapor')
            if props_v is not None:
                dT_v_dt = q_v / (props_v['density'] * props_v['cp'])
            else:
                dT_v_dt = 0.0
        else:
            dT_v_dt = 0.0
            
        return dT_l_dt, dT_v_dt 